# -*- coding: utf-8 -*-
import datetime, time
from typing import List, Set, Dict

from core import threadLocal
from core.clock import Clock
from core.context import Context
from core.dataClasses import R, TacticResultData, SecurityGetter
from core.enums import Period, ThreadName
from core.tactic import Tactic
from event.event_manager import EventManager, Event
from infrastructure.util import dateutils
from infrastructure.util.dateutils import DateFormat
from core import logger


class SecuritySelection:
    """
    选股
    """
    def __init__(self):
        self.event_manager = EventManager.getInstance(threadLocal.get(ThreadName.EVENT_MANAGER.value))
        # InventedTrader.getInstance(self.name)(fund=10000, t=1)
        self.context = Context.getInstance(threadLocal.get(ThreadName.CONTEXT.value))
        self.clock = Clock.getInstance(threadLocal.get(ThreadName.CLOCK.value))
        self.tactic: Tactic = None
        self.securities: Set[str] = []
        self.startDate: datetime.date = None

        r = R.sucess(f"初始化选股模块成功")
        logger.info(r)

    def addSecurities(self, securities: List[str]):
        self.securities.extend(securities)
        r = R.sucess(f"添加股票代码{len(self.securities)}个成功")
        logger.info(r)

    def setTactic(self, tactic: Tactic):
        self.tactic = tactic
        r = R.sucess(f"设置策略[{tactic.classInfo}]成功")
        logger.info(r)
        return self

    def setStartDate(self, date: str, dateFormat: DateFormat = DateFormat.Y_m_d):
        if date is None:
            self.startDate = datetime.datetime.today()
        else:
            self.startDate = datetime.datetime.strptime(date, dateFormat.value)
        r = R.sucess(f"设置回测开始时间[{date}]成功")
        logger.info(r)

    def getStartDate(self):
        if self.startDate is None:
            self.startDate = datetime.datetime.strptime(dateutils.getLastestTradingDate(), DateFormat.Y_m_d.value)
        return self.startDate

    def _run(self) -> List[str]:
        assert self.tactic is not None
        assert len(self.securities) > 0
        for security in self.securities:
            r = R.sucess(f"开始检测股票{security}")
            logger.info(r)
            self.clock.setCurrentDateTime(date=dateutils.getLastestTradingDate(dateFormat=DateFormat.Y_m_d), dateFormat=DateFormat.Y_m_d, lastestFinishedUnit=True, period=self.context.period)
            self.context.set_Security(security)
            while not self.clock.isEnd(unit=self.context.period):
                self.tactic.compute()
                self.clock.flow(unit=self.context.period)
        result = []
        if self.tactic.notes is not None and len(self.tactic.notes) > 0:
            result = [s for s in self.tactic.notes]
        return result

    def start(self):
        self._run()


class SecuritySelectionMultiprocessing(SecuritySelection):
    def __init__(self):
        super(SecuritySelectionMultiprocessing, self).__init__()
        self.getter = None
        self.lock = None
        self.shareDict: Dict[List[str]] = None
        self.period: Period = None

    def setGetter(self, getter: SecurityGetter):
        self.getter = getter
        return self

    def _run(self):
        assert self.tactic is not None
        assert self.getter is not None
        assert self.lock is not None
        assert self.shareDict is not None
        assert self.period is not None
        generator = self.getter(self.lock, self.shareDict, self.period)
        while(True):
            try:
                security = next(generator)
            except:
                generator = self.getter(self.lock, self.shareDict, self.period)
                security = next(generator, None)
            if security is None or security == " ":
                # time.sleep(1)
                continue
            r = R.sucess(f"开始检测股票{security}")
            logger.info(r)
            self.clock.setCurrentDateTime(date=dateutils.getLastestTradingDate(dateFormat=DateFormat.Y_m_d),
                                          dateFormat=DateFormat.Y_m_d, lastestFinishedUnit=True,
                                          period=self.context.period)
            self.context.set_Security(security)
            while not self.clock.isEnd(unit=self.context.period):
                self.tactic.compute()
                self.clock.flow(unit=self.context.period)


class SecuritySelectionMultiThread(SecuritySelection):
    def __init__(self):
        super(SecuritySelectionMultiThread, self).__init__()
        self.getter = None

    def setGetter(self, getter: SecurityGetter):
        self.getter = getter
        return self

    def _run(self):
        assert self.tactic is not None
        assert self.getter is not None
        try:
            generator = self.getter()
        except:
            pass
        while(True):
            try:
                securityInfo: TacticResultData = next(generator)
            except:
                try:
                    generator = self.getter()
                except:
                    pass
                # time.sleep(1)
                continue
            if securityInfo is None:
                continue
            r = R.sucess(f"开始检测股票{securityInfo.security}")
            logger.info(r)
            self.context.set_Security(securityInfo.security)
            self.tactic.compute()
